package comtgl.demo.test;

import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;

public class TokenBucketLimiter {
    // 秒级令牌桶
    private final TokenBucket secondBucket;
    // 分钟级令牌桶
    private final TokenBucket minuteBucket;

    public TokenBucketLimiter(int secondRate, int minuteRate) {
        this.secondBucket = new TokenBucket(secondRate, 1000);
        this.minuteBucket = new TokenBucket(minuteRate, 60 * 1000);
    }

    /**
     * 尝试获取令牌
     * @return true表示获取成功，false表示被限流
     */
    public boolean tryAcquire() {
        return secondBucket.tryAcquire() && minuteBucket.tryAcquire();
    }

    private static class TokenBucket {
        // 令牌桶容量
        private final int capacity;
        // 令牌产生速率(毫秒)
        private final long rate;
        // 当前令牌数量
        private final AtomicInteger tokens;
        // 上次填充时间
        private final AtomicLong lastRefillTime;

        public TokenBucket(int capacity, long periodMs) {
            this.capacity = capacity;
            this.rate = periodMs / capacity;
            this.tokens = new AtomicInteger(capacity);
            this.lastRefillTime = new AtomicLong(System.currentTimeMillis());
        }

        /**
         * 尝试获取令牌
         * @return true表示获取成功，false表示被限流
         */
        public boolean tryAcquire() {
            refill();
            int currentTokens = tokens.get();
            if (currentTokens <= 0) {
                return false;
            }
            while (true) {
                currentTokens = tokens.get();
                if (currentTokens <= 0) {
                    return false;
                }
                if (tokens.compareAndSet(currentTokens, currentTokens - 1)) {
                    return true;
                }
            }
        }

        /**
         * 填充令牌
         */
        private void refill() {
            long now = System.currentTimeMillis();
            long lastTime = lastRefillTime.get();

            // 计算应该填充的令牌数
            long elapsed = now - lastTime;
            int generatedTokens = (int) (elapsed / rate);

            if (generatedTokens <= 0) {
                return;
            }

            // 填充令牌（使用CAS保证线程安全）
            while (true) {
                int currentTokens = tokens.get();
                int newTokens = Math.min(capacity, currentTokens + generatedTokens);
                if (lastRefillTime.compareAndSet(lastTime, now)) {
                    tokens.set(newTokens);
                    break;
                }
                // 重新获取最后填充时间
                lastTime = lastRefillTime.get();
                elapsed = now - lastTime;
                generatedTokens = (int) (elapsed / rate);
                if (generatedTokens <= 0) {
                    break;
                }
            }
        }
    }
}